"""
============================================================
🧾 File: logger_manager.py
🧠 Description:
    本模块实现训练过程中的日志与可视化管理器 (LoggerManager)，
    用于统一处理以下功能：
        ✅ 控制台 + 文件日志输出
        ✅ TensorBoard 可视化支持
        ✅ 自动记录训练与验证指标
        ✅ 按时间生成唯一日志文件

✨ Features:
    - 自动创建日志目录
    - 按时间命名日志文件
    - 支持 TensorBoard 输出
    - 统一格式的日志输出格式 (时间 | 等级 | 消息)

📌 Usage:
    >>> logger = LoggerManager(train_config)
    >>> logger.log_train(epoch, step, loss, metrics)
    >>> logger.log_val(epoch, metrics)
    >>> logger.close()

🚀 Author: kkli
📅 Date: 2025-10-30
============================================================
"""

import os
import logging
from torch.utils.tensorboard import SummaryWriter
from datetime import datetime


class LoggerManager:
    """
    日志与可视化管理类，用于管理训练和验证阶段的日志输出。
    支持文件日志、控制台输出以及 TensorBoard 可视化。

    Args:
        config (TrainConfig): 训练配置类，需包含以下字段：
            - log_dir (str): 日志保存路径
            - use_tensorboard (bool): 是否启用 TensorBoard
    """

    def __init__(self, config):
        """初始化 LoggerManager"""
        self.config = config
        os.makedirs(config.log_dir, exist_ok=True)

        # === 设置日志文件路径 ===
        log_file = os.path.join(
            config.log_dir,
            f"train_{datetime.now().strftime('%Y%m%d_%H%M%S')}.log"
        )

        # === 初始化 logging 配置 ===
        logging.basicConfig(
            level=logging.INFO,
            format="%(asctime)s | %(levelname)s | %(message)s",
            handlers=[
                logging.FileHandler(log_file, encoding="utf-8"),
                logging.StreamHandler()
            ]
        )

        # === 初始化 TensorBoard ===
        self.writer = SummaryWriter(log_dir=config.log_dir) if config.use_tensorboard else None

        logging.info("🧩 LoggerManager 初始化完成")

    # -----------------------------------------------------
    def log_train(self, epoch: int, step: int, loss: float, metrics: dict):
        """
        记录训练阶段日志与 TensorBoard 可视化数据。

        Args:
            epoch (int): 当前训练轮次
            step (int): 当前训练步数
            loss (float): 当前批次的损失值
            metrics (dict): 其他训练指标，如 {"accuracy": 0.92, "f1": 0.87}

        Example:
            >>> logger.log_train(epoch=1, step=100, loss=0.23, metrics={"acc": 0.94})
        """
        msg = f"[Train] Epoch {epoch} Step {step} | Loss: {loss:.4f}"
        for k, v in metrics.items():
            msg += f" | {k}: {v:.4f}"
        logging.info(msg)

        # 写入 TensorBoard 可视化
        if self.writer:
            self.writer.add_scalar("train/loss", loss, step)
            for k, v in metrics.items():
                self.writer.add_scalar(f"train/{k}", v, step)

    # -----------------------------------------------------
    def log_val(self, epoch: int, metrics: dict):
        """
        记录验证阶段日志与 TensorBoard 数据。

        Args:
            epoch (int): 当前 epoch
            metrics (dict): 验证指标字典，如 {"loss": 0.31, "iou": 0.78}

        Example:
            >>> logger.log_val(epoch=2, metrics={"loss": 0.31, "iou": 0.78})
        """
        msg = f"[Val] Epoch {epoch} | " + " | ".join([f"{k}: {v:.4f}" for k, v in metrics.items()])
        logging.info(msg)

        # 写入 TensorBoard
        if self.writer:
            for k, v in metrics.items():
                self.writer.add_scalar(f"val/{k}", v, epoch)

    # -----------------------------------------------------
    def close(self):
        """
        关闭日志资源，释放 TensorBoard 句柄。
        """
        if self.writer:
            self.writer.close()
            logging.info("🧹 TensorBoard writer 已关闭")


def cross_entropy_loss(p, q):
    """
    交叉熵损失函数（Cross Entropy Loss）

    数学公式：
        L = -∑ y_i * log(p_i)

    或使用 LaTeX 表示：
        .. math::
            L = -\sum_i y_i \log(p_i)

    Args:
        p (Tensor): 预测概率分布
        q (Tensor): 实际标签分布（one-hot）
    """
    import torch
    return -torch.sum(q * torch.log(p))
